# encoding=utf-8
import os
import pickle
import tempfile
import cv2 as cv
from abc import abstractmethod, ABC
from nnunet.training.model_restore import *
from nnunet.training.network_training import network_trainer
from nnunet.postprocessing.connected_components import load_postprocessing
from seg_common.PredictService import PredictFactory
import numpy as np
import SimpleITK as sitk
from seg_system.Segmentation.service.Nnunet.NnunetPreProcess import NnunetPreProcess

from seg_system.application_config import ApplicationConfig

"""关于nnunet的部署

Developer: tacom
Date: 2023.4.18

tips:
    这里的代码大多于框架代码二次封装而来,
    此框架对代码作了很多泛化,对于此次的权重没有用到的部分,会进行一些删减


plans.pkl
    num_stages
    num_modalities
    modalities
    normalization_schemes
    dataset_properties
    list_of_npz_files
    original_spacings
    original_sizes
    preprocessed_data_folder
    num_classes
    all_classes
    base_num_features
    use_mask_for_norm
    keep_only_largest_region
    min_region_size_per_class
    min_size_per_class
    transpose_forward
    transpose_backward
    data_identifier
    plans_per_stage
    preprocessor_name


"""


class NnunetBase(PredictFactory):
    def before(self, obj, **kwargs):
        return obj

    def forward(self, obj, **kwargs):
        return self.trainer.predict_preprocessed_data_return_seg_and_softmax(
            obj, do_mirroring=True, mirror_axes=self.trainer.data_aug_params['mirror_axes'],
            use_sliding_window=True, step_size=0.5, use_gaussian=True, all_in_gpu=False,
            mixed_precision=True
        )

    def after(self, obj, **kwargs):
        transpose_forward = self.trainer.plans.get('transpose_forward')
        if transpose_forward is not None:
            transpose_backward = self.trainer.plans.get('transpose_backward')
            softmax = obj[1].transpose([0] + [i + 1 for i in transpose_backward])
            return softmax
        return obj[1]

    def save(self, obj, path, name, **kwargs):
        # index_list = kwargs.get('index_list')
        properties_list = kwargs.get("properties")
        from_path = kwargs.get('from_path')
        tmp_path = kwargs.get('tmp_path')

        # 这里nnunet的框架代码,会存储3D的数据,这里和2D 到 3D的情况不一样
        # 不是一次性处理的,所以直接产生临时文件,或者尝试在内存中直接进行转换
        for idx in range(len(name)):
            each_name = name[idx]
            tmp_save_path = tempfile.NamedTemporaryFile(suffix='.nii.gz')

            each_properties = properties_list[idx]

            # 运算方式比较别致
            each_obj = obj[idx * 2: (idx + 1) * 2]
            seg_old_spacing = each_obj.argmax(0)  # 二分类处理方法

            seg_resized_itk = sitk.GetImageFromArray(seg_old_spacing.astype(np.uint8))
            seg_resized_itk.SetSpacing(each_properties['itk_spacing'])
            seg_resized_itk.SetOrigin(each_properties['itk_origin'])
            seg_resized_itk.SetDirection(each_properties['itk_direction'])
            sitk.WriteImage(seg_resized_itk, tmp_save_path.name)

            # 转换后的3D数据, 需要转换到2D进行存储
            convert_mat = NnunetPreProcess.convert2png(tmp_save_path.name)

            # 同时需要存储分割结果与原图合并的图像
            ori_mat = cv.imread(os.path.join(from_path, each_name))

            # 恢复数据的相对位置
            copy_mask = np.zeros(ori_mat.shape[:2], np.uint8)
            h_s, h_e = each_properties['crop_bbox'][1]
            w_s, w_e = each_properties['crop_bbox'][2]
            copy_mask[h_s: h_e, w_s: w_e] = convert_mat

            convert_mat = copy_mask
            cv.imwrite(os.path.join(path, each_name), convert_mat * 255)

            no_zero_idx = np.nonzero(convert_mat)
            ori_mat[no_zero_idx[0], no_zero_idx[1], :] = \
                ori_mat[no_zero_idx[0], no_zero_idx[1], :] * 0.5 + np.array([0, 212, 255]) * 0.5
            cv.imwrite(os.path.join(tmp_path, each_name), ori_mat)

    def __init__(self,
                 model_path: str,
                 fold_num: int = 4,
                 check_point_name: str = 'model_best',
                 mixed_precision: bool = True):
        # 公共的模型初始化参数
        self.model_path = model_path
        self.fold_num = fold_num
        self.check_point_name = check_point_name
        self.mixed_precision = mixed_precision

        # 代码看到的,暂时先添加上去
        self.force_separate_z = None
        self.interpolation_order = 1
        self.interpolation_order_z = 0

        # 初始化
        self.trainer: network_trainer = None  # 这个trainer实际上就是predictor
        self.params = None  # 这个实际上就是模型的参数,不知道为什么独立开来
        self.checkpoint_init()

    def checkpoint_init(self):
        """

        nnunet的模型初始化较为复杂,独立开来

        :param
        :return
        """
        self.trainer, self.params = load_model_and_checkpoint_files(
            self.model_path, self.fold_num,
            mixed_precision=self.mixed_precision,
            checkpoint_name=self.check_point_name
        )
        self.trainer.load_checkpoint_ram(self.params[0], False)
        self.trainer.network.eval()


class NnunetBaseMulti(NnunetBase):
    def __init__(self, model_path: str, fold_num: int = 4, check_point_name: str = 'model_best',
                 mixed_precision: bool = True, color_list: list = None):
        # 公共的模型初始化参数
        super().__init__(model_path, fold_num, check_point_name, mixed_precision)
        self.plans_path = os.path.join(model_path, "plans.pkl")
        self.class_list: list = []
        self.color_list = color_list

        self.load_multi_info()

    def load_multi_info(self):
        # 可能是代码已经废弃,直接读取描述json已经获取不到这个分类模型的具体数据了
        # 继续改为plans.pkl文件
        # for_which_classes, min_valid_obj_size = load_postprocessing(self.pre_post_json)
        # self.class_list = for_which_classes
        with open(self.plans_path, 'rb') as f:
            plans = pickle.load(f)
            self.class_list = plans['all_classes']

    def save(self, obj, path, name, **kwargs):
        # index_list = kwargs.get('index_list')
        properties_list = kwargs.get("properties")
        from_path = kwargs.get('from_path')
        tmp_path = kwargs.get('tmp_path')

        # 这里nnunet的框架代码,会存储3D的数据,这里和2D 到 3D的情况不一样
        # 不是一次性处理的,所以直接产生临时文件,或者尝试在内存中直接进行转换
        for idx in range(len(name)):
            each_name = name[idx]
            tmp_save_path = tempfile.NamedTemporaryFile(suffix='.nii.gz')

            each_properties = properties_list[idx]

            # 多分类处理方法, 类别1, 类别2, ..., 类别背景
            seg_old_spacing_final = np.zeros(obj.shape[1:])
            for i, c in enumerate(self.class_list):
                seg_old_spacing_final[obj[i + 1] > 0.5] = c
            seg_old_spacing = seg_old_spacing_final

            seg_resized_itk = sitk.GetImageFromArray(seg_old_spacing.astype(np.uint8))
            seg_resized_itk.SetSpacing(each_properties['itk_spacing'])
            seg_resized_itk.SetOrigin(each_properties['itk_origin'])
            seg_resized_itk.SetDirection(each_properties['itk_direction'])
            sitk.WriteImage(seg_resized_itk, tmp_save_path.name)

            # 转换后的3D数据, 需要转换到2D进行存储
            convert_mat = NnunetPreProcess.convert2png(tmp_save_path.name)

            # 同时需要存储分割结果与原图合并的图像
            ori_mat = cv.imread(os.path.join(from_path, each_name))

            # 恢复数据的相对位置
            copy_mask = np.zeros(ori_mat.shape[:2], np.uint8)
            h_s, h_e = each_properties['crop_bbox'][1]
            w_s, w_e = each_properties['crop_bbox'][2]
            copy_mask[h_s: h_e, w_s: w_e] = convert_mat

            convert_mat = copy_mask
            # TODO 业务系统冲突, 对于非零的地方统一255处理
            save_mat = np.zeros(convert_mat.shape, dtype=np.uint8)
            save_mat[convert_mat > 0] = 255
            cv.imwrite(os.path.join(path, each_name), save_mat)

            # 对原图进行染色
            for cls_index in self.class_list:
                each_cls_idx = np.where(convert_mat == cls_index)
                if (self.color_list is not None) and (len(self.color_list) == len(self.class_list)):
                    rand_color = np.array(self.color_list[cls_index - 1])
                else:
                    rand_color = np.array([0, (212 + 20 * cls_index) % 255, 255])
                ori_mat[each_cls_idx[0], each_cls_idx[1], :] = \
                    ori_mat[each_cls_idx[0], each_cls_idx[1], :] * 0.5 + rand_color * 0.5
            cv.imwrite(os.path.join(tmp_path, each_name), ori_mat)


class NnunetBig(NnunetBase):

    def __new__(cls, *args, **kwargs):
        if not hasattr(cls, '_inst_big'):
            cls._inst_big = super(NnunetBig, cls).__new__(cls)
        return cls._inst_big

    def __init__(self):
        super().__init__(os.path.join(ApplicationConfig.PathConfig.SYSTEM_RESOURCE,
                                      ApplicationConfig.NetConfig.SEGMENTATION_DICT[
                                          ApplicationConfig.NetConfig.SEGMENTATION_USE_NNUNET_BIG]))


class NnunetSmall(NnunetBase):
    def __new__(cls, *args, **kwargs):
        if not hasattr(cls, '_inst_small'):
            cls._inst_small = super(NnunetSmall, cls).__new__(cls)
        return cls._inst_small

    def __init__(self):
        super().__init__(os.path.join(ApplicationConfig.PathConfig.SYSTEM_RESOURCE,
                                      ApplicationConfig.NetConfig.SEGMENTATION_DICT[
                                          ApplicationConfig.NetConfig.SEGMENTATION_USE_NNUNET_SMALL]))


class NnunetFaz(NnunetBase):
    def __new__(cls, *args, **kwargs):
        if not hasattr(cls, '_inst_faz'):
            cls._inst_faz = super(NnunetFaz, cls).__new__(cls)
        return cls._inst_faz

    def __init__(self):
        super().__init__(os.path.join(ApplicationConfig.PathConfig.SYSTEM_RESOURCE,
                                      ApplicationConfig.NetConfig.SEGMENTATION_DICT[
                                          ApplicationConfig.NetConfig.SEGMENTATION_USE_NNUNET_FAZ]))


class NnunetIntersection(NnunetBaseMulti):
    def __new__(cls, *args, **kwargs):
        if not hasattr(cls, '_inst_intersection'):
            cls._inst_intersection = super(NnunetIntersection, cls).__new__(cls)
        return cls._inst_intersection

    def __init__(self):
        super().__init__(os.path.join(ApplicationConfig.PathConfig.SYSTEM_RESOURCE,
                                      ApplicationConfig.NetConfig.SEGMENTATION_DICT[
                                          ApplicationConfig.NetConfig.SEGMENTATION_USE_NNUNET_INTERSECTION]),
                         color_list=[[0, 212, 255], [0, 252, 255], [255, 0, 0], [0, 17, 255]])


if __name__ == '__main__':
    pass
